from vertex import Vertex
from edge import Edge
import numpy
from OpenGL.GL import *

WIREFRAME = 0
SOLID = 1

class Curtain():
    pointMass = 1.0
    gravity = numpy.array([0.0, -10.0, 0.0], dtype='float32') # for starters

    def __init__(self, numRows, numCols, spacing, centre=numpy.zeros(3, dtype='float32'), mode=WIREFRAME):
        self.rows = numRows
        self.cols = numCols
        self.spacing = spacing
        self.mode = mode
        
        self.centre = centre
        self.min = centre - numpy.array([0.5 * numCols * spacing, 0.0, 0.0])
        self.max = centre + numpy.array([0.5 * numCols * spacing, numRows * spacing, 0.0])        
        
        self.reset()
        
        # set shader sources
        self.vshader_src = """
            #version 110
            
            uniform mat4 mvTransform;
            uniform mat4 pmvTransform;            
            
            varying vec3 N;
            varying vec3 v;
            void main(void)  
            {     
               v = vec3(mvTransform * gl_Vertex);       
               N = normalize(gl_Normal);
               gl_Position = pmvTransform * gl_Vertex;  
            }
            """
            
        self.fshader_src = """
            #version 110
           
            varying vec3 N;
            varying vec3 v;    
            uniform vec4 lightPos;
            uniform vec4 ambient;
            uniform vec4 diffuse;
            uniform vec4 specular;
            uniform float shininess;
            uniform vec4 color;
            
            void main (void)  
            {  
               vec3 L = normalize(lightPos.xyz - v);   
               vec3 E = normalize(-v); // we are in Eye Coordinates, so EyePos is (0,0,0)  
               vec3 R = normalize(-reflect(L,N));  
             
               //calculate Ambient Term:  
               vec4 Iamb = ambient;    
            
               //calculate Diffuse Term:  
               vec4 Idiff = diffuse * max(dot(N,L), 0.0);
               Idiff = clamp(Idiff, 0.0, 1.0);     
               
               // calculate Specular Term:
               vec4 Ispec = specular * pow(max(dot(R,E),0.0),0.3*shininess);
               Ispec = clamp(Ispec, 0.0, 1.0);               
               
               // write Total Color:  
               gl_FragColor = color * (Idiff + Iamb + Ispec); // vec4(N, 1.0);
            }
            """
            
        self.color = numpy.array([0.85, 0.25, 0.25, 1.0], dtype='float32')
        self.glStatus = False
        
    def update(self, dT):
        # apply gravity to each vertex
        for i in range(self.rows):
            for j in range(self.cols):
                self.vertices[i][j].applyForce(self.gravity, dT)
        
        for edge in self.edges:
            edge.solveContraints(dT)
            
        
        # reset min / max before updating positions
        self.min = numpy.zeros(3)
        self.max = numpy.zeros(3)
        
        # finally update the position of each vertex
        for i in range(self.rows):
            for j in range(self.cols):
                curVertex = self.vertices[i][j] 
                curVertex.updatePos(dT)
                
                # track the current min / max for the bounding box
                self.min = numpy.minimum(self.min, curVertex.pos)
                self.max = numpy.maximum(self.max, curVertex.pos)
        
    def boundingBox(self):
        return self.min, self.max
            
    def initialized(self):
        return self.glStatus        
            
    def initializeGL(self):
        # create mesh buffers for vertices / edges mesh
        self.drawVertices = self._getDrawVertices()
        self.drawEdges = self._getDrawEdges()       
        self.normals = self._getNormals()   
    
        # compile and link shader
        vshader_id = glCreateShader(GL_VERTEX_SHADER)
        glShaderSource(vshader_id, self.vshader_src)
        glCompileShader(vshader_id)
        print("curtain vertex shader compile status: "+
                str(glGetShaderiv(vshader_id, GL_COMPILE_STATUS)))
        print(glGetShaderInfoLog(vshader_id))

        fshader_id = glCreateShader(GL_FRAGMENT_SHADER)
        glShaderSource(fshader_id, self.fshader_src)
        glCompileShader(fshader_id)
        print("curtain fragment shader compile status: "+
                str(glGetShaderiv(fshader_id, GL_COMPILE_STATUS)))
        print(glGetShaderInfoLog(fshader_id))

        self.shader = glCreateProgram()
        glAttachShader(self.shader, vshader_id)
        glAttachShader(self.shader, fshader_id)
        glLinkProgram(self.shader)

        self.mvLocation = glGetUniformLocation(self.shader, "mvTransform")
        self.pmvLocation = glGetUniformLocation(self.shader, "pmvTransform")
        self.colorLocation = glGetUniformLocation(self.shader, "color")
        
        self.lightPosLocation = glGetUniformLocation(self.shader, "lightPos")
        self.lightAmbientLocation = glGetUniformLocation(self.shader, "ambient")
        self.lightDiffuseLocation = glGetUniformLocation(self.shader, "diffuse")
        self.lightSpecularLocation = glGetUniformLocation(self.shader, "specular")
        self.lightShininessLocation = glGetUniformLocation(self.shader, "shininess")

        # then set up VBO stuff
        self.meshbuffers = glGenBuffers(3)
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[0])
        glBufferData(GL_ARRAY_BUFFER, self.drawVertices, GL_STATIC_DRAW)
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[1])
        glBufferData(GL_ARRAY_BUFFER, self.normals, GL_STATIC_DRAW)
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.meshbuffers[2])
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.drawEdges, GL_STATIC_DRAW)        
        
        # remember that initialization is complete
        self.glStatus = True

    def _getNormals(self):
        normals = list()
        for i in range(self.rows):
            for j in range(self.cols):
                curVertex = self.vertices[i][j]

                t1 = numpy.zeros(3)
                t2 = numpy.zeros(3)

                if i == 0: # no top neighbour element
                    t1 = self.vertices[i+1][j].getPos() - curVertex.getPos()                   
 
                elif i == self.rows - 1: # no bottom neighbour element
                    t1 = curVertex.getPos() - self.vertices[i-1][j].getPos()

                else:
                    t1 = self.vertices[i+1][j].getPos() - self.vertices[i-1][j].getPos()

                if j == 0: # no left neighbour element
                    t2 = self.vertices[i][j+1].getPos() - curVertex.getPos()                   

                elif j == self.cols - 1: # no right neighbour element
                    t2 = curVertex.getPos() - self.vertices[i][j-1].getPos()

                else:
                    t2 = self.vertices[i][j+1].getPos() - self.vertices[i][j-1].getPos()

                curNormal = numpy.cross(t1, t2)
                normals.extend([curNormal[0], curNormal[1], curNormal[2]])

        return numpy.array(normals, dtype='float32')                
    
    def _getDrawVertices(self):
        drawVerticesList = list()
        for i in range(self.rows):
            for j in range(self.cols):
                curVertex = self.vertices[i][j]
                drawVerticesList.extend(curVertex.getPos())
                drawVerticesList.extend([1.0])
                
        return numpy.array(drawVerticesList, dtype='float32')
    
    def _getDrawEdges(self):
        if(self.mode == WIREFRAME):
            drawEdgeList = list()
            for i in range(self.rows):
                for j in range(self.cols): 
                    curVertex = (i * self.cols) + j
                    if i < self.rows - 1:
                        # edge to neighbour below
                        drawEdgeList.extend([curVertex, curVertex+self.cols])
                        
                    if j < self.cols - 1:
                        # edge to neighbour to right
                        drawEdgeList.extend([curVertex, curVertex+1])
                        
                        # edge to neigbhour below and to the right
                    if i < self.rows - 1 and j < self.cols - 1:
                        drawEdgeList.extend([curVertex, curVertex+self.cols+1])
                        
            return numpy.array(drawEdgeList, dtype='uint32')
        elif (self.mode == SOLID):
            drawEdgeList = list()
            for i in range(self.rows):
                for j in range(self.cols):
                    curVertex = (i * self.cols) + j
                    if i != 0 and j != self.cols - 1:
                        #draw upper triangle
                        drawEdgeList.extend([curVertex, curVertex+1, curVertex - self.cols])
                    if i != self.rows - 1 and j !=0:
                        #draw lower triangle
                        drawEdgeList.extend([curVertex, curVertex-1, curVertex + self.cols])  
            return numpy.array(drawEdgeList, dtype='uint32')
        
    def paintGL(self, mvTransform, pmvTransform, light):
        # update drawVertices with the new positions of the vertices        
        self.drawVertices = self._getDrawVertices()
        self.drawEdges = self._getDrawEdges()
        self.normals = self._getNormals()
        
        glUseProgram(self.shader)
        glUniformMatrix4fv(self.mvLocation, 1, True, mvTransform.flatten())
        glUniformMatrix4fv(self.pmvLocation, 1, True, pmvTransform.flatten())
        glUniform4fv(self.colorLocation, 1, self.color)
        glUniform4fv(self.lightPosLocation, 1, light.getPos())
        glUniform4fv(self.lightAmbientLocation, 1, light.getAmbient())
        glUniform4fv(self.lightDiffuseLocation, 1, light.getDiffuse())
        glUniform4fv(self.lightSpecularLocation, 1, light.getSpecular())
        glUniform1f(self.lightShininessLocation, light.getShininess())
        
        # vertex data
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[0])
        glBufferData(GL_ARRAY_BUFFER, self.drawVertices.flatten(), GL_STATIC_DRAW)
        glVertexPointer(4, GL_FLOAT, 0, None)
        glEnableClientState(GL_VERTEX_ARRAY)        

        # normal data
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[1])
        glBufferData(GL_ARRAY_BUFFER, self.normals, GL_STATIC_DRAW)        
        glNormalPointer(GL_FLOAT, 0, None)
        glEnableClientState(GL_NORMAL_ARRAY)
        
        # surface data (edges / triangles)
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.meshbuffers[2])
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.drawEdges, GL_STATIC_DRAW)
        if(self.mode == WIREFRAME):
            glDrawElements(GL_LINES, len(self.drawEdges), GL_UNSIGNED_INT, None)
        elif(self.mode == SOLID):
            glDrawElements(GL_TRIANGLES, len(self.drawEdges), GL_UNSIGNED_INT, None)        
                
        glDisableClientState(GL_NORMAL_ARRAY)
        glDisableClientState(GL_VERTEX_ARRAY)
    
    def changeMode(self, mode):
        self.mode = mode
        self.drawEdges = self._getDrawEdges()
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.meshbuffers[1])
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.drawEdges, GL_STATIC_DRAW)
        
    def reset(self):
        self.vertices = list()        
        for i in range(self.rows):
            newRow= list()
            
            for j in range(self.cols):                
                fixedVertex = False
                if i == 0:
                    fixedVertex = True                
                
                # centre the curtain around self.centre
                xPos = (j * self.spacing) - (0.5 * self.cols * self.spacing)
                yPos = ((self.rows - i) * self.spacing)
                
                zPos = 0.0
                halfRows = self.rows / 2.0
                if i > halfRows:
                    zPos = i - halfRows
                    
                getPos = numpy.array([xPos, yPos, zPos])                
                newVertex = Vertex(self.pointMass, getPos, fixed=fixedVertex)
                newRow.append(newVertex)
                
            self.vertices.append(newRow)
                
        # now traverse the 2D array of vertices and create edges between adjacents
        self.edges = list()
        for i in range(self.rows):
            for j in range(self.cols):
                
                v1 = self.vertices[i][j]
                if i < self.rows-1:
                    v2 = self.vertices[i+1][j]    
                    newEdge = Edge(v1, v2)
                    self.edges.append(newEdge)
                    
                if j < self.cols-1:
                    v2 = self.vertices[i][j+1]
                    newEdge = Edge(v1, v2)
                    self.edges.append(newEdge)
        
        
